import pygame
import sys
import argparse
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from env.ball_battle_env import BallBattleEnv
from agents.dqn_agent import DQNAgent
from stable_baselines3 import DQN
import torch
import numpy as np

def play_with_human_control():
    """人类玩家控制游戏"""
    print("Starting game with human control...")
    print("Controls: Mouse to aim, Left click to shoot, W/S to move forward/backward")
    
    # 创建环境
    env = BallBattleEnv(render_mode="human")
    
    # 初始化Pygame
    pygame.init()
    clock = pygame.time.Clock()
    
    # 游戏主循环
    observation, info = env.reset()
    running = True
    
    while running:
        # 获取鼠标位置
        mouse_x, mouse_y = pygame.mouse.get_pos()
        
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                
        # 设置玩家球球朝向（指向鼠标位置）
        if env.player_ball and env.player_ball.alive:
            from utils.math_utils import angle_between_points
            player_pos = (env.player_ball.x, env.player_ball.y)
            mouse_pos = (mouse_x, mouse_y)
            angle = angle_between_points(player_pos, mouse_pos)
            env.player_ball.set_direction(angle)
                
        # 获取键盘输入
        keys = pygame.key.get_pressed()
        action = 0  # 默认停止
        
        # 解析按键为动作
        if keys[pygame.K_w]:  # 前进
            action = 1
        elif keys[pygame.K_s]:  # 后退
            action = 2
        # 鼠标左键射击
        elif pygame.mouse.get_pressed()[0]:  # 左键按下
            action = 5
        # 移除了A/D转向和空格键射击，因为现在使用鼠标控制
            
        # 执行动作
        observation, reward, terminated, truncated, info = env.step(action)
        
        # 渲染环境
        env.render()
        
        # 控制帧率
        clock.tick(60)
        
        # 检查游戏结束
        if terminated or truncated:
            print(f"Game Over! Team Scores: {info['team_scores']}")
            observation, info = env.reset()
            
    # 关闭环境
    env.close()
    pygame.quit()
    sys.exit()

def play_with_ai_control(model_path=None, use_sb3=False):
    """AI控制游戏"""
    print("Starting game with AI control...")
    
    # 创建环境
    env = BallBattleEnv(render_mode="human")
    
    # 加载模型
    if use_sb3:
        # 使用Stable-Baselines3模型
        model = DQN.load(model_path if model_path else "models/sb3_dqn_model")
        # 为所有AI球球设置模型
        env.ai_models = [model]  # 简化处理，所有AI球球使用同一个模型
    else:
        # 使用自定义DQN模型
        agent = DQNAgent(env.action_space, env.observation_space)
        if model_path:
            agent.load(model_path)
            # 为所有AI球球设置模型
            env.ai_models = [agent]  # 简化处理，所有AI球球使用同一个模型
        else:
            print("Warning: No model path provided, using random actions")
            env.ai_models = []  # 空列表表示使用随机动作
    
    # 初始化Pygame
    pygame.init()
    clock = pygame.time.Clock()
    
    # 游戏主循环
    observation, info = env.reset()
    running = True
    
    while running:
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                
        # AI决策
        if use_sb3:
            action, _states = model.predict(observation, deterministic=True)
        else:
            action = agent.predict(observation)
            
        # 执行动作
        observation, reward, terminated, truncated, info = env.step(action)
        
        # 渲染环境
        env.render()
        
        # 控制帧率
        clock.tick(60)
        
        # 检查游戏结束
        if terminated or truncated:
            print(f"Game Over! Team Scores: {info['team_scores']}")
            observation, info = env.reset()
            
    # 关闭环境
    env.close()
    pygame.quit()
    sys.exit()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Play Ball Battle Game")
    parser.add_argument("--mode", choices=["human", "ai"], default="human",
                        help="Choose game mode: human control or AI control (default: human)")
    parser.add_argument("--model", type=str, help="Path to trained model")
    parser.add_argument("--sb3", action="store_true", help="Use Stable-Baselines3 model")
    args = parser.parse_args()
    
    if args.mode == "human":
        play_with_human_control()
    else:
        play_with_ai_control(args.model, args.sb3)

if __name__ == "__main__":
    main()